import ustruct
from machine import Pin, I2S
import time
import _thread
from oledprint import getI2Cstatue, threaded_print, OLEDPrinter


class Recorder:
    def __init__(self, i2s_config, button_pin, led_pin, wav_filename="audio.wav"):
        self.wav_filename = wav_filename
        self.sample_rate = i2s_config['rate']
        self.bits_per_sample = i2s_config['bits']
        self.channels = 1
        self.buffer_size = 1024
        self.recording = False
        self.wav_file = None
        self.exit_program = False  # 添加退出标志

        # 初始化 I2S 接口
        self.i2s = I2S(0, **i2s_config)

        # 初始化 LED
        self.led = Pin(led_pin, Pin.OUT)

        # 初始化按钮
        self.button = Pin(button_pin, Pin.IN, Pin.PULL_UP)

        # 初始化 OLED 打印
        w = getI2Cstatue()
        if w:
            self.oled_printer = OLEDPrinter(w)
            global print  # 声明 print 为全局变量

            def print(*args, **kwargs):
                threaded_print(self.oled_printer, *args, **kwargs)
        print("Recorder initialized, ready to start.")

    def create_wav_header(self, data_size):
        """创建 WAV 文件头部"""
        byte_rate = self.sample_rate * self.channels * self.bits_per_sample // 8
        block_align = self.channels * self.bits_per_sample // 8
        return b'RIFF' + \
               ustruct.pack('<I', 36 + data_size) + \
               b'WAVE' + \
               b'fmt ' + \
               ustruct.pack('<IHHIIHH', 16, 1, self.channels, self.sample_rate,
                            byte_rate, block_align, self.bits_per_sample) + \
               b'data' + \
               ustruct.pack('<I', data_size)

    def start_record(self):
        """开始录音"""
        data_size = 3600 * self.sample_rate * self.bits_per_sample // 8 * self.channels  # 最大录制1小时
        self.wav_file = open(self.wav_filename, "wb")
        wav_header = self.create_wav_header(data_size)
        self.wav_file.write(wav_header)
        self.recording = True
        print("start record...")
        self.led.value(1)

    def stop_record(self):
        """停止录音并退出程序"""
        self.recording = False
        if self.wav_file:
            self.wav_file.close()
            self.wav_file = None
        print(f"record over, saved as {self.wav_filename}")
        self.led.value(0)
        self.exit_program = True  # 设置退出标志为 True

    def handle_record(self):
        """录音处理逻辑"""
        buffer = bytearray(self.buffer_size)
        while self.recording:
            num_bytes = self.i2s.readinto(buffer)
            if num_bytes > 0 and self.wav_file:
                self.wav_file.write(buffer[:num_bytes])
            time.sleep(0.01)  # 减少 CPU 占用率

    def button_listener(self):
        """按钮监听线程"""
        last_state = self.button.value()
        while not self.exit_program:  # 添加检查退出程序的标志
            current_state = self.button.value()
            if current_state != last_state:
                if not current_state:  # 按钮按下
                    if not self.recording:
                        self.start_record()
                else:  # 按钮释放
                    if self.recording:
                        self.stop_record()
                last_state = current_state
            time.sleep(0.05)  # 防止按钮抖动

    def run(self):
        """启动录音器"""
        print("准备就绪，按下按钮开始录音，释放按钮停止录音。")
        # 启动按钮监听线程
        _thread.start_new_thread(self.button_listener, ())
        # 主线程处理录音数据
        try:
            while not self.exit_program:  # 当 exit_program 为 False 时继续运行
                if self.recording:
                    self.handle_record()
                time.sleep(0.1)
        except KeyboardInterrupt:
            self.stop_record()
            self.i2s.deinit()
            print("程序已停止。")
        finally:
            self.exit_program = True  # 确保设置退出标志，让按钮监听线程退出



# I2S 配置
i2s_config = {
    'sck': Pin(27),  # I2S SCK 引脚
    'ws': Pin(25),   # I2S WS 引脚
    'sd': Pin(26),   # I2S SD 引脚
    'mode': I2S.RX,  # 接收模式
    'bits': 16,      # 位宽
    'format': I2S.MONO,  # 单声道
    'rate': 16000,   # 采样率
    'ibuf': 20000    # 缓冲区大小
}

# 创建并运行录音器
recorder = Recorder(i2s_config=i2s_config, button_pin=18, led_pin=2)
if __name__ == "__main__":
    recorder.run()